#include <stdio.h>
#include <stdlib.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>

int AVFrame2JPG(AVFrame *pfrm, char *jpgname) {

    AVFormatContext *pfmt_ctx = NULL;
    AVStream *pstm = NULL;
    AVCodecParameters *pcodecpar = NULL;
    AVPacket packet;
    AVCodecContext *pcodec_ctx = NULL;
    AVCodec *pcodec = NULL;
    int ret = 0;
    int width = pfrm->width;
    int height = pfrm->height;

    ret = avformat_alloc_output_context2(&pfmt_ctx, NULL, "mjpeg", jpgname);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }   

    if ((ret = avio_open2(&(pfmt_ctx->pb), jpgname, AVIO_FLAG_WRITE, NULL, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avio_open2 ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    pstm = avformat_new_stream(pfmt_ctx, NULL);
    if (pstm == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_new_stream ERROR in AVFrame2JPG \n");
        return AVERROR(ENOMEM);
    }

    pcodecpar = pstm->codecpar;
    pcodecpar->height = height;
    pcodecpar->width = width;
    pcodecpar->codec_id = pfmt_ctx->oformat->video_codec;
    pcodecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    pcodecpar->format = AV_PIX_FMT_YUVJ420P;

    pcodec = avcodec_find_encoder(pcodecpar->codec_id);
    if (pcodec == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder ERROR in AVFrame2JPG \n");
        return AVERROR(EINVAL);
    }
    //printf("codec name: %s, long name: %s \n", pcodec->name, pcodec->long_name);

    pcodec_ctx = avcodec_alloc_context3(pcodec);
    if (pcodec_ctx == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR in main \n");
        return AVERROR(ENOMEM);
    }

    ret = avcodec_parameters_to_context(pcodec_ctx, pstm->codecpar);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    pcodec_ctx->time_base.den = 25;
    pcodec_ctx->time_base.num = 1;

    //** change the quality of jpg
    pcodec_ctx->qcompress = 20.0/100.0;
    pcodec_ctx->qmin = 2;
    pcodec_ctx->qmax = 52;
    pcodec_ctx->max_qdiff = 5;

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    ret = avformat_write_header(pfmt_ctx, NULL);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_write_header ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    ret = av_new_packet(&packet, 3 * width * height);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_new_packet ERROR in AVFrame2JPG \n");
        return -1;
    }

    ret = avcodec_send_frame(pcodec_ctx, pfrm);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_frame ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    while (ret >= 0) {
        ret = avcodec_receive_packet(pcodec_ctx, &packet);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF ) {
            break;
        } else if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_packet ERROR %s in AVFrame2JPG \n", av_err2str(ret));
            return AVERROR(ret);
        }
        
        if ((ret = av_write_frame(pfmt_ctx, &packet)) < 0) {
            av_log(NULL, AV_LOG_ERROR, "av_write_frame ERROR %s in AVFrame2JPG \n", av_err2str(ret));
            return AVERROR(ret);
        }
        av_packet_unref(&packet);
        av_write_trailer(pfmt_ctx);
    }

    avcodec_close(pcodec_ctx);
    avformat_free_context(pfmt_ctx);
    avio_close(pfmt_ctx->pb);

    return 0;
}

int main(int argc, char **argv) {

    if (argc != 4){
        fprintf(stdout, "Usage: ./yuv2jpg ./xx.yuv width height\n");
        fprintf(stderr, "You must input yuv file! \n");
        return -1;
    }

    char *yuv_name = argv[1];
    int width = atoi(argv[2]);
    int height = atoi(argv[3]);
    uint8_t *buffer = NULL;
    FILE *fp = NULL;
    AVFrame *pfrm = NULL;
    int count = 0; //frame count
    int frame_size = 0;
    char outname[100];
    int ret = 0;

    fp = fopen(yuv_name, "rb");
    if (fp == NULL) {
        av_log(NULL, AV_LOG_ERROR, "fopen ERROR in main\n");
        return -1;
    }

    pfrm = av_frame_alloc();
    if (pfrm == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR in main \n");
        return AVERROR(ENOMEM);
    }
    
    frame_size = av_image_get_buffer_size(AV_PIX_FMT_YUVJ420P, width, height, 1);
    if (frame_size < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_image_get_buffer_size ERROR in main \n ");
        return -1;
    }
    

    buffer = (uint8_t *) av_mallocz(frame_size);
    if (buffer == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_mallocz ERROR in main \n");
        return -1;
    }

    pfrm->width = width;
    pfrm->height = height;
    pfrm->format = AV_PIX_FMT_YUVJ420P;

    while (fread(buffer, 1, width * height * 3/2, fp) == width * height * 3/2) {
        pfrm->data[0] = buffer;
        pfrm->data[1] = (buffer + width * height);
        pfrm->data[2] = buffer + width * height * 5/4;

        ret = av_image_fill_arrays(pfrm->data, pfrm->linesize, buffer, AV_PIX_FMT_YUV420P, width, height, 1);
        if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "av_image_fill_arrays ERROR in main \n");
            return AVERROR(ret);
        }

        count++;
        sprintf(outname, "%d.jpg", count);
        fprintf(stdout, "Output file name : %s \n", outname);
        ret = AVFrame2JPG(pfrm, outname);
        if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "AVFrame2JPG ERROR in main \n");
            break;
        }
    }

    fclose(fp);
    av_freep(&buffer);
    av_frame_free(&pfrm);

    return 0;
}